"""
    Title

    author: wxz
    date: 2021-12-21
    github: https://github.com/xinzwang
"""
import glob
import hashlib
import os
from multiprocessing import Pool
import xml.etree.ElementTree as ET

import numpy as np
from pathlib import Path
from tqdm import tqdm

CACHE_VERSION = '1.0.0'
NUM_THREADS = min(8, max(1, os.cpu_count() - 1))  # number of multiprocessing threads


class RAW_Normal(object):
    def __init__(self, path='../datasets/RAW_Normal', img_size=640, batch_size=16):
        self.path = path

        img_file = []


def voc_xml_label(path=None):
    assert path is not None, print("[ERROR]: Please input label path")
    # load label paths
    label_files = []
    for p in path if isinstance(path, list) else [path]:
        p = Path(p)
        if p.is_dir():
            label_files += glob.glob(str(p / '**' / '*.*'), recursive=True)
        elif p.is_file():  # train.txt  test.txt  val.txt
            with open(p) as t:
                t = t.read().strip().splitlines()
                parent = str(p.parent) + os.sep
                for x in t:
                    if x.startswith('./'):
                        label_files += [x.replace(('./', parent))]
                    else:
                        label_files += [x]
        else:
            raise Exception(f'{p} dose not exist')
    label_files = sorted(x.replace('/', os.sep) for x in label_files if x.split('.')[-1].lower() in ['xml'])
    assert label_files, print('[ERROR] No label was found')
    # cache labels
    cache_path = (p if p.is_file() else Path(label_files[0]).parent).with_suffix('cache')
    try:
        cache, exists = np.load(cache_path), True
        assert cache['version'] == CACHE_VERSION
        assert cache['hash'] == get_hash(label_files)
    except:
        print('\nCache dose note exist or has expired.')
        print('Trying to create new cache file ...')
        cache, exists = cache_xml_label(label_files, cache_path), False
    # display cache messages
    nf, nm, ne, nc, n = cache.pop('results')
    if exists:
        d = f"Scanning '{cache_path}' images and labels... {nf} found, {nm} missing, {ne} empty, {nc} corrupted"
        tqdm(None, desc=d, total=n, initial=n)  # display cache results
        if cache['msgs']:
            print(f"[INFO]: {cache['msgs']}")
    # phrase caceh
    [cache.pop(k) for k in ('hash', 'version', 'msgs')]


def cache_xml_label(label_files=None, path=None):
    # cache dataset labels, check images and read shapes
    x = {}
    nf, nm, ne, nc, msgs = 0, 0, 0, 0, []  # number missing, found, empty, corrupt, messages
    desc = f"Scanning '{path.parent / path.stem}' images and labels..."
    with Pool(NUM_THREADS) as pool:
        pbar = tqdm(pool.imap(None, label_files))
        for im_file, l, shape, segment, nm_f, nf_f, ne_f, nc_f, msg in pbar:
            nm += nm_f
            nf += nf_f
            ne += ne_f
            nc += nc_f
            if im_file:
                x[im_file] = [l, shape, segment]
            if msg:
                msgs.append(msg)
            pbar.desc = f"{desc}{nf} found, {nm} missing, {ne} empty, {nc} corrupted"
    pbar.close()
    if msgs:
        print(f"[INFO]: {msgs}")
    if nf == 0:
        print(f"[WARNING]: No labels found in {path}.")
    x['hash'] = get_hash(label_files)
    x['results'] = nf, nm, ne, nc, len(label_files)
    x['msgs'] = msgs
    x['version'] = CACHE_VERSION
    try:
        np.save(path, x)
        path.with_suffix('.cache.npy').rename('path')
        print(f"[NEWS]: New cache created: {path}")
    except Exception as e:
        print(f"[WARNING]: Save cache error: {e}")
    return x


def verify_xml_label(args):
    # verify label file. read label from xml files
    label_file = args
    nm, nf, ne, nc, msg, segments = 0, 0, 0, 0, '', []  # number (missing, found, empty, corrupt), message, segments
    try:
        # pharse VOC xml label file
        if os.path.isfile(label_file):
            nf = 1  # label found
            tree = ET.parse(label_file)
            root = tree.getroot()
            objs = root.findall('object')
            width = int(root.find('size').find('width').text)
            height = int(root.find('size').find('height').text)
            depth = int(root.find('size').find('depth').text)
            l = []
            for obj in objs:
                name = obj.find('name').text
                xmin = int(obj.find('bndbox').find('xmin').text)
                ymin = int(obj.find('bndbox').find('ymin').text)
                xmax = int(obj.find('bndbox').find('xmax').text)
                ymax = int(obj.find('bndbox').find('ymax').text)

            nl = len(l)



        else:
            nm = 1  # label missing
            l = np.zeros((0, 5), dtype=np.float32)  # TODO:check this
    except Exception as e:
        nc = 1  # corrupt
        msg = f'[WARNING]: {label_file}: ignoring corrupt image/label: {e}'
        return [None, None, None, None, nm, nf, ne, nc, msg]


def get_hash(paths):
    # Returns a single hash value of a list of paths (files or dirs)
    size = sum(os.path.getsize(p) for p in paths if os.path.exists(p))  # sizes
    h = hashlib.md5(str(size).encode())  # hash sizes
    h.update(''.join(paths).encode())  # hash paths
    return h.hexdigest()  # return hash

